{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# {func}`~tvm._ffi.registry.register_func` 注册全局函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mSignature:\u001b[0m \u001b[0mtvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_ffi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mregister_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moverride\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mSource:\u001b[0m   \n",
      "\u001b[0;32mdef\u001b[0m \u001b[0mregister_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moverride\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;34m\"\"\"Register global function\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    Parameters\u001b[0m\n",
      "\u001b[0;34m    ----------\u001b[0m\n",
      "\u001b[0;34m    func_name : str or function\u001b[0m\n",
      "\u001b[0;34m        The function name\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    f : function, optional\u001b[0m\n",
      "\u001b[0;34m        The function to be registered.\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    override: boolean optional\u001b[0m\n",
      "\u001b[0;34m        Whether override existing entry.\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    Returns\u001b[0m\n",
      "\u001b[0;34m    -------\u001b[0m\n",
      "\u001b[0;34m    fregister : function\u001b[0m\n",
      "\u001b[0;34m        Register function if f is not specified.\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    Examples\u001b[0m\n",
      "\u001b[0;34m    --------\u001b[0m\n",
      "\u001b[0;34m    The following code registers my_packed_func as global function.\u001b[0m\n",
      "\u001b[0;34m    Note that we simply get it back from global function table to invoke\u001b[0m\n",
      "\u001b[0;34m    it from python side. However, we can also invoke the same function\u001b[0m\n",
      "\u001b[0;34m    from C++ backend, or in the compiled TVM code.\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    .. code-block:: python\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m      targs = (10, 10.0, \"hello\")\u001b[0m\n",
      "\u001b[0;34m      @tvm.register_func\u001b[0m\n",
      "\u001b[0;34m      def my_packed_func(*args):\u001b[0m\n",
      "\u001b[0;34m          assert(tuple(args) == targs)\u001b[0m\n",
      "\u001b[0;34m          return 10\u001b[0m\n",
      "\u001b[0;34m      # Get it out from global function table\u001b[0m\n",
      "\u001b[0;34m      f = tvm.get_global_func(\"my_packed_func\")\u001b[0m\n",
      "\u001b[0;34m      assert isinstance(f, tvm.PackedFunc)\u001b[0m\n",
      "\u001b[0;34m      y = f(*targs)\u001b[0m\n",
      "\u001b[0;34m      assert y == 10\u001b[0m\n",
      "\u001b[0;34m    \"\"\"\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;32mif\u001b[0m \u001b[0mcallable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0mf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfunc_name\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0mfunc_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"expect string function name\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mioverride\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mctypes\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mc_int\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moverride\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;32mdef\u001b[0m \u001b[0mregister\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmyf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;34m\"\"\"internal register function\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmyf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mPackedFuncBase\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m            \u001b[0mmyf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconvert_to_tvm_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmyf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0mcheck_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_LIB\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTVMFuncRegisterGlobal\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mc_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmyf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mioverride\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;32mreturn\u001b[0m \u001b[0mmyf\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;32mif\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;32mreturn\u001b[0m \u001b[0mregister\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;32mreturn\u001b[0m \u001b[0mregister\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFile:\u001b[0m      /media/pc/data/lxw/ai/tvm/python/tvm/_ffi/registry.py\n",
      "\u001b[0;31mType:\u001b[0m      function"
     ]
    }
   ],
   "source": [
    "import tvm._ffi\n",
    "\n",
    "tvm._ffi.register_func??"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{func}`~tvm._ffi.registry.register_func` 用于注册全局函数。\n",
    "\n",
    "示例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.runtime.packed_func import PackedFunc\n",
    "targs = (10, 10.0, \"hello\")\n",
    "@tvm._ffi.register_func\n",
    "def my_packed_func(*args):\n",
    "    assert(tuple(args) == targs)\n",
    "    return 10\n",
    "# 从全局函数表中获取它\n",
    "f = tvm.get_global_func(\"my_packed_func\")\n",
    "assert isinstance(f, PackedFunc)\n",
    "y = f(*targs)\n",
    "assert y == 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
